from cnocr import CnOcr
import time
import re
import os


def extract_text():
    def get_image_files_in_folder(folder_path):
        image_files = []
        for file in os.listdir(folder_path):
            if file.endswith('.JPG') or file.endswith('.png') or file.endswith('.jpeg') or file.endswith(
                    '.bmp') or file.endswith('.jpg') or file.endswith('.PNG') or file.endswith(
                    '.JPEG') or file.endswith('.BMP'):
                image_files.append(os.path.join(folder_path, file))

        # 按照文件名中的数字部分进行排序
        image_files.sort(key=lambda f: int(re.search(r'(\d+)', os.path.basename(f)).group()))
        return image_files

    start_time = time.time()
    folder_path = 'output'
    image_files = get_image_files_in_folder(folder_path)
    ocr = CnOcr(name='densenet-s')
    ocr_ = CnOcr(det_model_name='en_PP-OCRv3_det', rec_model_name='en_PP-OCRv3')
    length = len(image_files)
    address_flag = 0
    address_result = ''
    count = -1
    result = []
    label_minzu = [
        "汉", "蒙古", "回", "藏", "维吾尔", "苗", "彝", "壮", "布依",
        "朝鲜", "满", "瑶", "白", "土家", "哈尼", "哈萨克", "傣", "黎",
        "傈僳", "佤", "畲", "高山", "拉祜", "水", "东乡", "纳西", "景颇",
        "柯尔克孜", "土", "达斡尔", "仫佬", "羌", "布朗", "撒拉", "毛南",
        "仡佬", "锡伯", "阿昌", "普米", "塔吉克", "怒", "乌孜别克", "俄罗斯",
        "鄂温克", "德昂", "保安", "裕固", "京", "塔塔尔", "独龙", "鄂伦春",
        "赫哲", "门巴", "珞巴", "基诺"
    ]
    for file_path in image_files:
        count = count + 1
        if count == len(image_files) - 1:
            ocr_data = ocr_.ocr(file_path)
            ocr_result = ocr_data[0]['text']
        else:
            ocr_data = ocr.ocr_for_single_line(file_path)
            ocr_result = ''.join([ocr_data['text']]).replace(' ', '')
        pattern = re.compile(r'[^\u4e00-\u9fff-\d]')  # 使用re.sub替换掉这些字符为空字符串
        ocr_result = re.sub(pattern, '', ocr_result)
        print("ocr_result:" + ocr_result)
        label = ['姓名', '性别', '民族', '出生', '住址', '公民身份号码']
        sex = ['男', '女']
        # 按照顺序判断姓名，性别，民族，出生
        if count < 3:
            ocr_result = re.sub(r'[^\u4e00-\u9fa5\d]', '', ocr_result)
            if count == 0:
                result.append('{}：{}'.format(label[count], ocr_result))
                continue
            if count == 1:
                for s in sex:
                    if s in ocr_result:
                        result.append('{}：{}'.format(label[count], s))
                        break
                continue
            if count == 2:
                for s in label_minzu:
                    if s in ocr_result:
                        result.append('{}：{}'.format(label[count], s))
                        break
                continue
        if count == 3:
            ocr_result = re.sub(r'[^0-9年月日]+', '', ocr_result)
            result.append('{}：{}'.format(label[count], ocr_result))
            continue
        # 先判断是不是住址区域，对住址数据进行操作
        if count == 4:
            if len(image_files) == 5:
                result.append('{}：{}'.format(label[count], ocr_result))
                continue
            else:
                address_result = ocr_result
                continue
        if count < len(image_files) - 1:
            address_result = address_result + ocr_result
        if count == len(image_files) - 2:
            result.append('{}：{}'.format(label[-2], address_result))
            continue
        if len(ocr_result) == 17:
            ocr_result = ocr_result + 'X'
        result.append('{}：{}'.format(label[-1], ocr_result))
    return result


def extract_formatted_text(current_date, uuids):
    def get_image_files_in_folder(folder_path):
        image_files = []
        for file in os.listdir(folder_path):
            if file.lower().endswith(('.jpg', '.png', '.jpeg', '.bmp')):
                image_files.append(os.path.join(folder_path, file))
        image_files.sort(key=lambda f: int(re.search(r'(\d+)', os.path.basename(f)).group()))
        return image_files

    start_time = time.time()
    folder_path = f'output/{current_date}/{uuids}'
    image_files = get_image_files_in_folder(folder_path)
    ocr = CnOcr(name='densenet-s')
    ocr_ = CnOcr(det_model_name='en_PP-OCRv3_det', rec_model_name='en_PP-OCRv3')
    label_minzu = [
        "汉", "蒙古", "回", "藏", "维吾尔", "苗", "彝", "壮", "布依",
        "朝鲜", "满", "瑶", "白", "土家", "哈尼", "哈萨克", "傣", "黎",
        "傈僳", "佤", "畲", "高山", "拉祜", "水", "东乡", "纳西", "景颇",
        "柯尔克孜", "土", "达斡尔", "仫佬", "羌", "布朗", "撒拉", "毛南",
        "仡佬", "锡伯", "阿昌", "普米", "塔吉克", "怒", "乌孜别克", "俄罗斯",
        "鄂温克", "德昂", "保安", "裕固", "京", "塔塔尔", "独龙", "鄂伦春",
        "赫哲", "门巴", "珞巴", "基诺"
    ]
    result = []
    sex = ['男', '女']

    # 处理前三个字段：姓名、性别、民族
    for i in range(3):
        if i >= len(image_files):
            break
        file_path = image_files[i]
        ocr_data = ocr.ocr_for_single_line(file_path)
        ocr_result = ''.join([ocr_data['text']]).replace(' ', '')
        ocr_result = re.sub(r'[^\u4e00-\u9fa5\d]', '', ocr_result)
        if i == 0:
            result.append({'姓名': ocr_result})
        elif i == 1:
            for s in sex:
                if s in ocr_result:
                    result.append({'性别': s})
                    break
        elif i == 2:
            for s in label_minzu:
                if s in ocr_result:
                    result.append({'民族': s})
                    break

    # 处理出生日期
    current_count = 3
    date_fragments = []
    while current_count < len(image_files) - 1:  # 保留身份证号位置
        file_path = image_files[current_count]
        ocr_data = ocr.ocr_for_single_line(file_path)
        ocr_result = ''.join([ocr_data['text']]).replace(' ', '')
        ocr_result = re.sub(r'[^\u4e00-\u9fa5\d]', '', ocr_result)
        print(f"处理出生日期{ocr_result},current_count={current_count}")

        # 改进条件判断，更严格地收集日期片段
        if re.search(r'年|月|日', ocr_result) or re.fullmatch(r'\d{4}', ocr_result):
            date_fragments.append(ocr_result)
            current_count += 1
        else:
            break

    # 终极解析策略改进
    date_info = {'year': None, 'month': None, 'day': None}

    # 第一步：通过关键字提取
    for fragment in date_fragments:
        # 提取年份
        year_match = re.search(r'(\d{4})年', fragment)
        if year_match and not date_info['year']:
            date_info['year'] = year_match.group(1)
        # 提取月份（允许无"月"字的情况）
        month_match = re.search(r'年(\d{1,2})(月|$)', fragment)  # 匹配"年XX"或"年XX月"
        if month_match and not date_info['month']:
            date_info['month'] = month_match.group(1)
        # 提取日期
        day_match = re.search(r'(\d{1,2})日', fragment)
        if day_match and not date_info['day']:
            date_info['day'] = day_match.group(1)

    # 第二步：数字序列解析补充
    merged_digits = ''.join(re.findall(r'\d+', ''.join(date_fragments)))
    if len(merged_digits) >= 4 and not date_info['year']:
        date_info['year'] = merged_digits[:4]
        merged_digits = merged_digits[4:]
    if len(merged_digits) >= 2 and not date_info['month']:
        possible_month = merged_digits[:2]
        if 1 <= int(possible_month) <= 12:
            date_info['month'] = possible_month
            merged_digits = merged_digits[2:]
    if len(merged_digits) >= 2 and not date_info['day']:
        possible_day = merged_digits[:2]
        if 1 <= int(possible_day) <= 31:
            date_info['day'] = possible_day

    # 第三步：最终验证和格式化
    try:
        if date_info['year'] and not (1900 < int(date_info['year']) < 2100):
            date_info['year'] = None
    except:
        date_info['year'] = None

    try:
        if date_info['month'] and not (1 <= int(date_info['month']) <= 12):
            date_info['month'] = None
    except:
        date_info['month'] = None

    try:
        if date_info['day'] and not (1 <= int(date_info['day']) <= 31):
            date_info['day'] = None
    except:
        date_info['day'] = None

    # 组合最终日期
    birth_date = []
    if date_info['year']:
        birth_date.append(f"{date_info['year']}年")
    if date_info['month']:
        birth_date.append(f"{date_info['month']}月")
    if date_info['day']:
        birth_date.append(f"{date_info['day']}日")
    result.append({'出生': ''.join(birth_date)})

    # ...（后续住址和身份证号处理保持不变）...

    # 处理住址（确保从current_count开始）
    address_result = ''
    while current_count < len(image_files) - 1:
        file_path = image_files[current_count]
        ocr_data = ocr.ocr_for_single_line(file_path)
        ocr_result = ''.join([ocr_data['text']]).replace(' ', '')
        ocr_result = re.sub(r'[^\u4e00-\u9fa5\d]', '', ocr_result)
        print(f"处理住址：{ocr_result},current_count={current_count}")
        address_result += ocr_result
        current_count += 1
    result.append({'住址': address_result})

    # 处理公民身份号码
    if len(image_files) > 0:
        file_path = image_files[-1]
        ocr_data = ocr_.ocr(file_path)
        ocr_result = ocr_data[0]['text']
        ocr_result = re.sub(r'[^\u4e00-\u9fa5\d]', '', ocr_result)
        print(f"处理公民身份号码:{ocr_result},current_count={current_count}")
        if len(ocr_result) == 17:
            ocr_result += 'X'
        result.append({'公民身份号码': ocr_result})

    return result, len(result)


# recognition.py

# recognition.py 中的 extract_formatted_text_back 函数优化
def extract_formatted_text_back(current_date, uuids):
    try:
        start = time.time()
        folder_path = os.path.join('output', str(current_date), str(uuids))
        if not os.path.exists(folder_path):
            return [], 0

        # 使用更适合中文的OCR模型
        ocr = CnOcr(name='densenet-s')
        # ocr = CnOcr(det_model_name='en_PP-OCRv3_det', rec_model_name='en_PP-OCRv3')
        ocr_ = CnOcr(det_model_name='en_PP-OCRv3_det', rec_model_name='en_PP-OCRv3')
        result = []

        # 签发机关处理（带格式清洗）
        issue_org_path = os.path.join(folder_path, 'sub_image_back_1.png')
        if os.path.exists(issue_org_path):
            # ocr_data = ocr.ocr(issue_org_path)
            # issue_text = ''.join([item['text'] for item in ocr_data])
            ocr_data = ocr.ocr_for_single_line(issue_org_path)
            issue_text = ''.join([ocr_data['text']]).replace(' ', '')
            # 强化清洗逻辑
            print("****************************************************")
            print(f"issue_text: {issue_text}")
            issue_text = re.sub(r'[^\u4e00-\u9fa5]|签发机关|[:：]', '', issue_text)
            result.append({'签发机关': issue_text.strip()})

        # 有效期限处理（带格式校验）
        valid_date_path = os.path.join(folder_path, 'sub_image_back_2.png')
        if os.path.exists(valid_date_path):
            ocr_data = ocr_.ocr(valid_date_path)
            date_text = ''.join([item['text'] for item in ocr_data])
            # 匹配标准日期格式
            match = re.search(r'(\d{4}年\d{2}月\d{2}日-\d{4}年\d{2}月\d{2}日)', date_text)
            if match:
                result.append({'有效期限': match.group(1)})
            else:
                # 尝试修复常见OCR错误
                date_text = re.sub(r'[Oo]', '0', date_text)  # 替换字母O
                date_text = re.sub(r'[^0-9年月日\-]', '', date_text)
                result.append({'有效期限': date_text})

        return result, len(result)
    except Exception as e:
        print(f"Error in extract_formatted_text_back: {str(e)}")
        return [], 0


def extract_formatted_text_back_old(current_date, uuids):
    def get_image_files_in_folder(folder_path):
        image_files = []
        for file in os.listdir(folder_path):
            if file.endswith('.JPG') or file.endswith('.png') or file.endswith('.jpeg') or file.endswith(
                    '.bmp') or file.endswith('.jpg') or file.endswith('.PNG') or file.endswith(
                    '.JPEG') or file.endswith('.BMP'):
                image_files.append(os.path.join(folder_path, file))

        # 按照文件名中的数字部分进行排序
        image_files.sort(key=lambda f: int(re.search(r'(\d+)', os.path.basename(f)).group()))
        return image_files

    try:
        start_time = time.time()
        folder_path = os.path.join('output', str(current_date), str(uuids))
        # 添加路径存在性校验（新增）
        if not os.path.exists(folder_path):
            raise FileNotFoundError(f"目标路径不存在：{folder_path}")

        # 添加路径调试信息（辅助排查）
        print(f"正在访问路径：{os.path.abspath(folder_path)}")
        image_files = get_image_files_in_folder(folder_path)
        ocr = CnOcr(name='densenet-s')
        ocr_ = CnOcr(det_model_name='en_PP-OCRv3_det', rec_model_name='en_PP-OCRv3')
        length = len(image_files)
        address_flag = 0
        address_result = ''
        count = -1
        result = []
        label = ['签发机关', '有效期限']

        for file_path in image_files:
            count = count + 1
            try:
                if count == len(image_files) - 1:
                    ocr_data = ocr_.ocr(file_path)
                    ocr_result = ocr_data[0]['text']
                else:
                    ocr_data = ocr.ocr_for_single_line(file_path)
                    ocr_result = ''.join([ocr_data['text']]).replace(' ', '')
                pattern = re.compile(r'[^\u4e00-\u9fff-\d]')  # 使用re.sub替换掉这些字符为空字符串
                ocr_result = re.sub(pattern, '', ocr_result)
                print("ocr_result:" + ocr_result)

                if '签发机关' in ocr_result:
                    cleaned_result = ocr_result.replace('签发机关', '')
                    result.append({label[0]: cleaned_result})
                elif '有效期限' in ocr_result:
                    cleaned_result = ocr_result.replace('有效期限', '')
                    result.append({label[1]: cleaned_result})
                print(f"result:{result}")
            except Exception as e:
                print(f"Error processing file {file_path}: {str(e)}")

        # print("进入到了这一步")
        return result, len(result)
    except Exception as e:
        print(f"Error in extract_formatted_text_back: {str(e)}")
        return [], 0